# resourceloader.py
import yaml
import cv2
import numpy as np
from PIL import Image
import os


class ResourcesLoader:
    def __init__(self, resources_dir="resources"):
        self.resources_dir = resources_dir
        self.resolution = None
        self.time = None
        self.victory_reward = None
        self.robot_vision = None
        self.infantry_hp = None
        self.hero_hp = None
        self.sentry_hp = None
        self.infantry_reborn_time = None
        self.hero_reborn_time = None
        self.sentry_reborn_time = None
        self.infantry_attack = None
        self.hero_attack = None
        self.sentry_attack = None
        self.infantry_attack_frequency = None
        self.hero_attack_frequency = None
        self.sentry_attack_frequency = None
        self.infantry_bullet_max = None
        self.hero_bullet_max = None
        self.sentry_bullet_max = None
        self.bullet_load_time = None
        self.attack_accuracy = {}  # dict for distances
        self.reward_center_occupy = None
        self.reward_kill = None

        # Map resources
        self.map_img = None
        self.map_walkable = None
        self.map_obstacles = None
        self.map_bonus = None
        self.map_red_ammo_zone = None
        self.map_blue_ammo_zone = None
        self.red_infantry_start = None
        self.red_hero_start = None
        self.red_sentry_start = None
        self.blue_infantry_start = None
        self.blue_hero_start = None
        self.blue_sentry_start = None
        self.map_real_width = None
        self.map_real_height = None
        self.map_scale = None

        # Robot resources
        self.robot_base = {}  # infantry, hero, sentry -> base mask
        self.robot_armor = {}  # -> armor mask
        self.robot_type_mark = {}  # -> type mark mask
        self.robot_real_width = None
        self.robot_real_height = None
        self.robot_scale = None

        self.load_params()
        self.load_map()
        self.load_robots()

    def load_params(self):
        with open(
            os.path.join(self.resources_dir, "params.yaml"), "r", encoding="utf-8"
        ) as f:
            params = yaml.safe_load(f)
        self.resolution = params["resolution"]
        self.time = params["time"]
        self.victory_reward = params["victory_reward"]
        self.robot_vision = params["robot_vision"]
        self.infantry_hp = params["infantry_hp"]
        self.hero_hp = params["hero_hp"]
        self.sentry_hp = params["sentry_hp"]
        self.infantry_reborn_time = params["infantry_reborn_time"]
        self.hero_reborn_time = params["hero_reborn_time"]
        self.sentry_reborn_time = params["sentry_reborn_time"]
        self.infantry_attack = params["infantry_attack"]
        self.hero_attack = params["hero_attack"]
        self.sentry_attack = params["sentry_attack"]
        self.infantry_attack_frequency = params["infantry_attack_frequency"]
        self.hero_attack_frequency = params["hero_attack_frequency"]
        self.sentry_attack_frequency = params["sentry_attack_frequency"]
        self.infantry_bullet_max = params["infantry_bullet_max"]
        self.hero_bullet_max = params["hero_bullet_max"]
        self.sentry_bullet_max = params["sentry_bullet_max"]
        self.bullet_load_time = params["bullet_load_time"]
        self.attack_accuracy = {
            1: params["attack_accuracy_1m"],
            2: params["attack_accuracy_2m"],
            3: params["attack_accuracy_3m"],
            5: params["attack_accuracy_5m"],
            8: params["attack_accuracy_8m"],
            20: params["attack_accuracy_20m"],
        }
        self.reward_center_occupy = params["reward_center_occupy"]
        self.reward_kill = params["reward_kill"]

    def load_map(self):
        with open(
            os.path.join(self.resources_dir, "map.yaml"), "r", encoding="utf-8"
        ) as f:
            map_config = yaml.safe_load(f)
        map_path = os.path.join(self.resources_dir, map_config["path"])
        self.map_real_width = map_config["real_width"]
        self.map_real_height = map_config["real_height"]

        # Load map image first without scaling
        map_img = cv2.imread(map_path)
        map_img = cv2.cvtColor(map_img, cv2.COLOR_BGR2RGB)
        original_h, original_w = map_img.shape[:2]

        # Calculate scale based on original dimensions
        self.map_scale = self.resolution / self.map_real_width

        # Parse colors BEFORE scaling
        white = np.array([255, 255, 255])
        black = np.array([0, 0, 0])
        red_inf = np.array([255, 0, 0])
        red_hero = np.array([255, 160, 160])
        red_sentry = np.array([255, 176, 176])
        blue_inf = np.array([0, 0, 255])
        blue_hero = np.array([160, 160, 255])
        blue_sentry = np.array([176, 176, 255])
        green = np.array([0, 255, 0])
        yellow = np.array([255, 255, 0])
        cyan = np.array([0, 255, 255])

        # Extract features from original image
        map_walkable = np.all(map_img == white, axis=-1)
        map_obstacles = np.all(map_img == black, axis=-1)
        map_bonus = np.all(map_img == green, axis=-1)
        map_red_ammo = np.all(map_img == yellow, axis=-1)
        map_blue_ammo = np.all(map_img == cyan, axis=-1)

        # Extract starting positions from original image
        red_infantry_start = np.argwhere(np.all(map_img == red_inf, axis=-1))
        red_hero_start = np.argwhere(np.all(map_img == red_hero, axis=-1))
        red_sentry_start = np.argwhere(np.all(map_img == red_sentry, axis=-1))
        blue_infantry_start = np.argwhere(np.all(map_img == blue_inf, axis=-1))
        blue_hero_start = np.argwhere(np.all(map_img == blue_hero, axis=-1))
        blue_sentry_start = np.argwhere(np.all(map_img == blue_sentry, axis=-1))

        # Resize everything to target resolution
        target_w = int(self.map_real_width * self.resolution)
        target_h = int(self.map_real_height * self.resolution)

        self.map_img = cv2.resize(map_img, (target_w, target_h))
        self.map_walkable = cv2.resize(
            map_walkable.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        self.map_obstacles = cv2.resize(
            map_obstacles.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        self.map_bonus = cv2.resize(
            map_bonus.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        self.map_red_ammo_zone = cv2.resize(
            map_red_ammo.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        self.map_blue_ammo_zone = cv2.resize(
            map_blue_ammo.astype(np.uint8), (target_w, target_h)
        ).astype(bool)

        # Scale positions according to new size
        pos_scale_x = target_w / original_w
        pos_scale_y = target_h / original_h
        if len(red_infantry_start) > 0:
            self.red_infantry_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in red_infantry_start
                ]
            )
        else:
            self.red_infantry_start = red_infantry_start
        if len(red_hero_start) > 0:
            self.red_hero_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in red_hero_start
                ]
            )
        else:
            self.red_hero_start = red_hero_start
        if len(red_sentry_start) > 0:
            self.red_sentry_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in red_sentry_start
                ]
            )
        else:
            self.red_sentry_start = red_sentry_start
        if len(blue_infantry_start) > 0:
            self.blue_infantry_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in blue_infantry_start
                ]
            )
        else:
            self.blue_infantry_start = blue_infantry_start
        if len(blue_hero_start) > 0:
            self.blue_hero_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in blue_hero_start
                ]
            )
        else:
            self.blue_hero_start = blue_hero_start
        if len(blue_sentry_start) > 0:
            self.blue_sentry_start = np.array(
                [
                    [int(p[1] * pos_scale_x), int(p[0] * pos_scale_y)]
                    for p in blue_sentry_start
                ]
            )
        else:
            self.blue_sentry_start = blue_sentry_start

        # For visualization: black/white + bonus + ammo zones
        vis_map = np.zeros_like(self.map_img)
        resized_obstacles = cv2.resize(
            map_obstacles.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        resized_bonus = cv2.resize(
            map_bonus.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        resized_red_ammo = cv2.resize(
            map_red_ammo.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        resized_blue_ammo = cv2.resize(
            map_blue_ammo.astype(np.uint8), (target_w, target_h)
        ).astype(bool)
        vis_map[resized_obstacles] = [0, 0, 0]
        vis_map[resized_bonus] = [0, 255, 0]
        vis_map[resized_red_ammo] = [255, 255, 0]
        vis_map[resized_blue_ammo] = [0, 255, 255]
        vis_map[
            ~resized_obstacles & ~resized_bonus & ~resized_red_ammo & ~resized_blue_ammo
        ] = [255, 255, 255]
        self.map_vis = vis_map

    def load_robots(self):
        with open(
            os.path.join(self.resources_dir, "robot.yaml"), "r", encoding="utf-8"
        ) as f:
            robot_config = yaml.safe_load(f)
        robot_path = os.path.join(self.resources_dir, robot_config["path"])
        self.robot_real_width = robot_config["real_width"]
        self.robot_real_height = robot_config.get(
            "real_height", robot_config["real_width"]
        )  # assume square if not specified
        self.robot_scale = self.resolution  # pixels per meter

        # Load robot image without scaling first
        robot_img = cv2.imread(robot_path)
        robot_img = cv2.cvtColor(robot_img, cv2.COLOR_BGR2RGB)
        original_h, original_w = robot_img.shape[:2]

        # Parse masks from original image
        black = np.array([0, 0, 0])
        green = np.array([0, 255, 0])
        magenta = np.array([255, 0, 255])

        base_mask = np.all(robot_img == black, axis=-1)
        armor_mask = np.all(robot_img == green, axis=-1)
        type_mask = np.all(robot_img == magenta, axis=-1)

        # Resize the parsed masks to target size based on real dimensions
        target_w = int(self.robot_real_width * self.robot_scale)
        target_h = int(self.robot_real_height * self.robot_scale)
        self.robot_base["general"] = cv2.resize(
            (base_mask * 255).astype(np.uint8), (target_w, target_h)
        )[:, :, np.newaxis]
        self.robot_armor["general"] = cv2.resize(
            (armor_mask * 255).astype(np.uint8), (target_w, target_h)
        )[:, :, np.newaxis]
        self.robot_type_mark["general"] = cv2.resize(
            (type_mask * 255).astype(np.uint8), (target_w, target_h)
        )[:, :, np.newaxis]

    def params_printer(self):
        """
        打印所有已解析的参数，用于检查参数是否正常加载
        """
        print("=== 参数检查 ===")
        print(f"分辨率: {self.resolution}")
        print(f"比赛时间: {self.time}")
        print(f"胜利奖励: {self.victory_reward}")
        print(f"机器人视野: {self.robot_vision}")

        print("\n=== 机器人血量 ===")
        print(f"步兵血量: {self.infantry_hp}")
        print(f"英雄血量: {self.hero_hp}")
        print(f"哨兵血量: {self.sentry_hp}")

        print("\n=== 机器人重生时间 ===")
        print(f"步兵重生时间: {self.infantry_reborn_time}")
        print(f"英雄重生时间: {self.hero_reborn_time}")
        print(f"哨兵重生时间: {self.sentry_reborn_time}")

        print("\n=== 机器人攻击属性 ===")
        print(f"步兵攻击力: {self.infantry_attack}")
        print(f"英雄攻击力: {self.hero_attack}")
        print(f"哨兵攻击力: {self.sentry_attack}")

        print("\n=== 机器人攻击频率 ===")
        print(f"步兵攻击频率: {self.infantry_attack_frequency}")
        print(f"英雄攻击频率: {self.hero_attack_frequency}")
        print(f"哨兵攻击频率: {self.sentry_attack_frequency}")

        print("\n=== 机器人弹药 ===")
        print(f"步兵最大弹药: {self.infantry_bullet_max}")
        print(f"英雄最大弹药: {self.hero_bullet_max}")
        print(f"哨兵最大弹药: {self.sentry_bullet_max}")
        print(f"装弹时间: {self.bullet_load_time}")

        print("\n=== 攻击精度 ===")
        for distance, accuracy in self.attack_accuracy.items():
            print(f"{distance}米精度: {accuracy}")

        print("\n=== 奖励设置 ===")
        print(f"占领中心奖励: {self.reward_center_occupy}")
        print(f"击杀奖励: {self.reward_kill}")

        print("\n=== 地图信息 ===")
        print(f"地图实际宽度: {self.map_real_width}")
        print(f"地图实际高度: {self.map_real_height}")
        print(f"地图缩放比例: {self.map_scale}")

        print("\n=== 机器人信息 ===")
        print(f"机器人实际宽度: {self.robot_real_width}")
        print(f"机器人实际高度: {self.robot_real_height}")
        print(f"机器人缩放比例: {self.robot_scale}")

        print("\n=== 起始位置 ===")
        print(f"红方步兵起始位置: {self.red_infantry_start}")
        print(f"红方英雄起始位置: {self.red_hero_start}")
        print(f"红方哨兵起始位置: {self.red_sentry_start}")
        print(f"蓝方步兵起始位置: {self.blue_infantry_start}")
        print(f"蓝方英雄起始位置: {self.blue_hero_start}")
        print(f"蓝方哨兵起始位置: {self.blue_sentry_start}")

    def show_render(self):
        """
        将解析出的机器人起始位置绘制在底图上并显示，用于检测位置是否正确
        """
        # 创建底图副本用于绘制
        render_img = self.map_img.copy()

        # 定义不同机器人的颜色
        colors = {
            "red_infantry": (255, 0, 0),  # 红色
            "red_hero": (255, 160, 160),  # 浅红色
            "red_sentry": (255, 176, 176),  # 更浅的红色
            "blue_infantry": (0, 0, 255),  # 蓝色
            "blue_hero": (160, 160, 255),  # 浅蓝色
            "blue_sentry": (176, 176, 255),  # 更浅的蓝色
        }

        # 绘制起始位置点
        positions = {
            "red_infantry": self.red_infantry_start,
            "red_hero": self.red_hero_start,
            "red_sentry": self.red_sentry_start,
            "blue_infantry": self.blue_infantry_start,
            "blue_hero": self.blue_hero_start,
            "blue_sentry": self.blue_sentry_start,
        }

        # 在图像上绘制每个位置的点
        for robot_type, pos_array in positions.items():
            if len(pos_array) > 0:
                color = colors[robot_type]
                for pos in pos_array:
                    # pos now [x, y]
                    x, y = pos[0], pos[1]
                    # Clamp for display
                    x = min(max(x, 0), render_img.shape[1] - 1)
                    y = min(max(y, 0), render_img.shape[0] - 1)
                    # 绘制一个稍大一点的圆点以便观察
                    cv2.circle(render_img, (int(x), int(y)), 3, color, -1)
                    # 添加文本标签
                    cv2.putText(
                        render_img,
                        robot_type.split("_")[1][:1],
                        (int(x) + 5, int(y) + 5),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.4,
                        color,
                        1,
                        cv2.LINE_AA,
                    )

        # 使用OpenCV显示图像
        # 转换回BGR格式用于OpenCV显示
        render_img_bgr = cv2.cvtColor(render_img, cv2.COLOR_RGB2BGR)
        cv2.imshow("Robot Starting Positions", render_img_bgr)
        cv2.waitKey(0)  # 等待按键
        cv2.destroyAllWindows()


if __name__ == "__main__":
    loader = ResourcesLoader()
    loader.params_printer()
    loader.show_render()
